import torch
import config
import numpy as np

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# 从配置文件获取维度信息
layer_num = config.LAYER_TOTAL_NUM
node_num = config.EDGE_NODE_NUM
user_num = config.USER_NUM

class RandomAgent:
    def __init__(self, env, device="cuda", random_seed=42):

        # 设备和随机种子
        self.device = torch.device(device if torch.cuda.is_available() else "cpu")
        torch.manual_seed(random_seed)

    def select_action(self, env, obs):
        # 准备用户位置信息
        # print("obs['tasks_layers'] shape:", np.shape(obs['tasks_layers']))
        # print("obs['nodes_layers'] shape:", np.shape(obs['nodes_layers']),"\n")
        each_use_loc = [
            [user[0], user[1], user[2]] for user in obs["users"]  # [id, x, y]
        ]

        # 准备边缘节点位置信息
        each_edge_loc = [
            [i, node[7], node[8]]  # [node_id, x, y]
            for i, node in enumerate(obs["nodes"])
        ]

        # 默认所有节点不可用
        available_actions = np.zeros(config.EDGE_NODE_NUM + 1, dtype=int)

        # 如果没有待处理任务列表，直接返回
        if not env.allexist_task_list:
            return available_actions

        # 获取当前任务
        current_task_idx = [
            item.usr_has_tsk for item in env.User if item.uid == env.next_uid_idx
        ][0][0]
        current_task = env.Task[current_task_idx]

        # 获取任务对应的用户
        task_user_map = env.get_task_user_mapping()
        matching_users = task_user_map.get(current_task.task_id, -1)

        if not matching_users:
            uid = -1
            print(
                "No user found for task_id:",
                current_task.task_id,
                current_task.assigned_node,
            )
            return config.EDGE_NODE_NUM, 0
        else:
            uid = matching_users

        # 使用env的pool方法筛选位置
        tem_next_can = env.pool(uid, each_use_loc, each_edge_loc)

        # 标记可用节点
        if current_task.reschedule_count <= 3:
            for node_index in tem_next_can:
                node = env.Edge[node_index]

                # 资源约束条件（注意使用task的实际值）
                is_resource_sufficient = (
                    node.container_number < config.node_max_container_number
                    and node.available_mem >= current_task.task_mem
                    and node.available_cpu >= current_task.task_cpu
                    and node.available_mem - current_task.task_mem
                    >= config.LIMITED_MEMORY * node.mem
                    and node.available_cpu - current_task.task_cpu
                    >= config.LIMITED_CPU * node.cpu
                )

                # 添加排除已分配节点的条件
                is_not_assigned_node = node_index != current_task.assigned_node

                if is_resource_sufficient and is_not_assigned_node:
                    available_actions[node_index] = 1
        else:
            # 如果重调度次数 > 3，所有边缘节点都标记为不可用
            available_actions = np.zeros(config.EDGE_NODE_NUM + 1, dtype=int)
            available_actions[config.EDGE_NODE_NUM] = 1
            return config.EDGE_NODE_NUM, 0

        # 找出可用动作的索引
        available_action_indices = np.where(available_actions == 1)[0]

        if len(available_action_indices) == 0:
            # 极端情况：没有可用动作，返回云节点
            action = config.EDGE_NODE_NUM
        else:
            action = np.random.choice(available_action_indices)

        return action, None #此处为原log_prob
